from functools import wraps

from flask import url_for, request, jsonify, Flask
from flask_restful import reqparse, abort, Api, Resource
from flask_jwt_extended import (
    create_access_token, create_refresh_token, jwt_required, get_jwt_identity, get_jwt, JWTManager,
    verify_jwt_in_request
)

"""
create_access_token() 用来创建 Token 令牌

get_jwt_identity() 用来根据令牌取得之前的 identity 信息

jwt_required() 这是一个装饰器，用来保护 flask 节点
"""
app = Flask(__name__)
# 密码
app.config['SECRET_KEY'] = '123'
# token可以放在headers、cookies中 https://mp.weixin.qq.com/s/hRT_VNhvIqJaaq6cDG0qhA
app.config["JWT_TOKEN_LOCATION"] = ["headers", "cookies", "json", "query_string"]
jwt = JWTManager(app)
api = Api(app)

# 自定义注解
# 判断是否含有is_admin,是否是admin，用于角色能力
def admin_required():
    def wrapper(fn):
        @wraps(fn)
        def decorator(*args, **kwargs):
            verify_jwt_in_request()
            claims = get_jwt()
            if claims["is_admin"]:
                return fn(*args, **kwargs)
            else:
                return jsonify(msg="Admins only!").json, 403

        return decorator

    return wrapper


class Login(Resource):

    # 生成
    def post(self):
        # 添加额外参数
        additional_claims = {"is_admin": False}
        access_token = create_access_token(identity='ssg', additional_claims=additional_claims)
        return jsonify({
            "code": "0",
            "message": "success",
            "data": {
                "access_token": access_token,
                "userid": 1001
            }
        })

    # optional=True 避免没有JWT报错，可以放回空等
    # 如果不存在 JWT，get_jwt()和 get_jwt_header() 将返回一个空字典。
    # get_jwt_identity(), current_user, 和get_current_user()将返回None。
    @jwt_required(optional=True)
    @admin_required()
    def get(self):
        """
        根据token 解析用户username
        Authorization: Bearer <access_token>
        """
        current_user = get_jwt_identity()
        return {
            "user": current_user
        }

    # 自定义超时返回
    @jwt.expired_token_loader
    def my_expired_token_callback(jwt_header, jwt_payload):
        """返回 flask Response 格式"""
        return jsonify(code="401", err="token 已过期"), 401

# 注册
api.add_resource(Login, '/api/v1/login')

if __name__ == '__main__':
    app.run()
